import torch
import torch.nn.functional as F

# 基于 https://github.com/pix2pixzero/pix2pix-zero 的代码
def noise_regularization(
    e_t,  # 当前步骤预测的噪声
    noise_pred_optimal,  # "最优"参考噪声（通常来自前向加噪过程）
    lambda_kl,  # KL散度正则化的权重系数
    lambda_ac,  # 自相关正则化的权重系数
    num_reg_steps,  # 正则化步骤的总数
    num_ac_rolls,  # 自相关损失计算的滚动次数
    generator=None  # 随机数生成器，用于可重复的随机性
):
    """
    对预测的噪声应用正则化，通过KL散度和自相关损失来优化噪声估计。
    这个过程通过梯度下降迭代地改进噪声预测。
    """
    # 执行多步正则化优化
    for _outer in range(num_reg_steps): # 通常为4步
        # KL散度正则化：确保预测噪声的分布与"最优"噪声分布相似
        if lambda_kl > 0:
            # 创建需要梯度的变量副本
            _var = torch.autograd.Variable(e_t.detach().clone(), requires_grad=True)
            # 计算KL散度损失
            l_kld = patchify_latents_kl_divergence(_var, noise_pred_optimal)
            # 反向传播计算梯度
            l_kld.backward()
            # 获取梯度并剪裁以防止梯度爆炸
            _grad = _var.grad.detach()
            _grad = torch.clip(_grad, -100, 100)
            # 使用梯度下降更新噪声估计
            e_t = e_t - lambda_kl * _grad
        
        # 自相关正则化：减少噪声中的周期性模式
        if lambda_ac > 0:
            # 多次计算自相关损失以获得更稳定的梯度估计
            for _inner in range(num_ac_rolls):
                # 创建需要梯度的变量副本
                _var = torch.autograd.Variable(e_t.detach().clone(), requires_grad=True)
                # 计算自相关损失
                l_ac = auto_corr_loss(_var, generator=generator)
                # 反向传播计算梯度
                l_ac.backward()
                # 获取梯度并平均（因为有多轮滚动）
                _grad = _var.grad.detach() / num_ac_rolls
                # 使用梯度下降更新噪声估计
                e_t = e_t - lambda_ac * _grad
        
        # 分离计算图，防止梯度累积
        e_t = e_t.detach()

    return e_t


# 基于 https://github.com/pix2pixzero/pix2pix-zero 的代码
def auto_corr_loss(
        x,  # 输入张量 [B, C, H, W]
        random_shift=True,  # 是否使用随机位移
        generator=None  # 随机数生成器
):
    """
    计算自相关损失，用于减少噪声中的周期性模式。
    通过在不同尺度和位移下计算噪声与其位移版本的相关性来实现。
    """
    B, C, H, W = x.shape
    assert B == 1  # 目前只支持批大小为1
    x = x.squeeze(0)  # 移除批次维度 [C, H, W]
    
    reg_loss = 0.0  # 初始化损失值
    
    # 对每个通道分别处理
    for ch_idx in range(x.shape[0]):
        # 提取当前通道的噪声 [1, 1, H, W]
        noise = x[ch_idx][None, None, :, :]
        
        # 在多尺度上计算自相关损失
        while True:
            # 随机选择位移量或使用固定位移
            if random_shift:
                roll_amount = torch.randint(0, noise.shape[2] // 2, (1,), generator=generator).item()
            else:
                roll_amount = 1
            
            # 计算水平方向的自相关
            reg_loss += (
                noise * torch.roll(noise, shifts=roll_amount, dims=2)  # 水平位移
            ).mean() ** 2  # 取平均后平方，使损失始终为正
            
            # 计算垂直方向的自相关
            reg_loss += (
                noise * torch.roll(noise, shifts=roll_amount, dims=3)  # 垂直位移
            ).mean() ** 2
            
            # 当特征图尺寸太小时停止多尺度处理
            if noise.shape[2] <= 8:
                break
            
            # 使用平均池化进行下采样，处理更粗糙的尺度
            noise = F.avg_pool2d(noise, kernel_size=2)
    
    return reg_loss


def patchify_latents_kl_divergence(x0, x1, patch_size=4, num_channels=4):
    """
    将潜在表示分割成小块，然后计算块之间的KL散度。
    这允许在局部区域而不是全局范围内比较分布。
    """
    
    def patchify_tensor(input_tensor):
        """
        将4D张量分割成重叠或非重叠的小块。
        使用unfold操作在空间维度上提取patch。
        """
        patches = (
            input_tensor.unfold(1, patch_size, patch_size)  # 在通道维度展开
            .unfold(2, patch_size, patch_size)  # 在高度维度展开
            .unfold(3, patch_size, patch_size)  # 在宽度维度展开
        )
        # 重新整形为 [num_patches, channels, patch_size, patch_size]
        patches = patches.contiguous().view(-1, num_channels, patch_size, patch_size)
        return patches

    # 将两个输入张量都分割成小块
    x0 = patchify_tensor(x0)
    x1 = patchify_tensor(x1)

    # 计算小块之间的KL散度并求和
    kl = latents_kl_divergence(x0, x1).sum()
    return kl


def latents_kl_divergence(x0, x1):
    """
    计算两个潜在表示集合之间的KL散度。
    基于每个集合的均值和方差进行计算。
    """
    EPSILON = 1e-6  # 小常数防止数值不稳定
    
    # 将空间维度展平 [N, C, H*W]
    x0 = x0.view(x0.shape[0], x0.shape[1], -1)
    x1 = x1.view(x1.shape[0], x1.shape[1], -1)
    
    # 计算均值和方差
    mu0 = x0.mean(dim=-1)  # 每个通道的均值
    mu1 = x1.mean(dim=-1)
    var0 = x0.var(dim=-1)  # 每个通道的方差
    var1 = x1.var(dim=-1)
    
    # 计算KL散度公式
    kl = (
        torch.log((var1 + EPSILON) / (var0 + EPSILON))  # 对数方差比
        + (var0 + (mu0 - mu1) ** 2) / (var1 + EPSILON)  # 均值差和方差项
        - 1  # 常数项
    )
    
    # 取绝对值并求和，确保损失为正
    kl = torch.abs(kl).sum(dim=-1)
    return kl


def inversion_step(
    pipe,  # 扩散模型管道
    z_t: torch.tensor,  # 时间步t的噪声潜在表示
    t: torch.tensor,  # 当前时间步
    prompt_embeds,  # 文本提示的嵌入表示
    added_cond_kwargs,  # 附加条件参数（如文本嵌入和时间ID）
    num_renoise_steps: int = 100,  # 重噪声步骤数
    first_step_max_timestep: int = 250,  # 区分第一步和其他步骤的时间步阈值
    generator=None,  # 随机数生成器
) -> torch.tensor:
    """
    执行扩散模型的反转步骤，使用噪声正则化来提高重建质量。
    这是DDIM反转的增强版本，包含多步噪声优化和平均。
    """
    extra_step_kwargs = {}
    
    # 根据当前时间步选择平均范围
    avg_range = pipe.cfg.average_first_step_range if t.item() < first_step_max_timestep else pipe.cfg.average_step_range
    # 根据当前时间步调整重噪声步骤数
    num_renoise_steps = min(pipe.cfg.max_num_renoise_steps_first_step, num_renoise_steps) if t.item() < first_step_max_timestep else num_renoise_steps

    nosie_pred_avg = None  # 用于存储平均噪声预测
    noise_pred_optimal = None  # 用于存储"最优"噪声参考
    
    # 通过前向加噪过程计算"真实"的下一时间步潜在表示
    z_tp1_forward = pipe.scheduler.add_noise(pipe.z_0, pipe.noise, t.view((1))).detach()
    
    # 初始化近似下一时间步的潜在表示
    approximated_z_tp1 = z_t.clone()
    
    # 多步重噪声优化循环
    for i in range(num_renoise_steps + 1):
        with torch.no_grad():  # 禁用梯度计算以提高效率
            # 如果启用了噪声正则化且是第一步，需要扩展批次大小
            if pipe.cfg.noise_regularization_num_reg_steps > 0 and i == 0:
                # 连接前向计算的潜在表示和当前近似值
                approximated_z_tp1 = torch.cat([z_tp1_forward, approximated_z_tp1])
                # 扩展提示嵌入
                prompt_embeds_in = torch.cat([prompt_embeds, prompt_embeds])
                
                # 扩展附加条件参数
                if added_cond_kwargs is not None:
                    added_cond_kwargs_in = {}
                    added_cond_kwargs_in['text_embeds'] = torch.cat([added_cond_kwargs['text_embeds'], added_cond_kwargs['text_embeds']])
                    added_cond_kwargs_in['time_ids'] = torch.cat([added_cond_kwargs['time_ids'], added_cond_kwargs['time_ids']])
                else:
                    added_cond_kwargs_in = None
            else:
                prompt_embeds_in = prompt_embeds
                added_cond_kwargs_in = added_cond_kwargs

            # 通过UNet获取噪声预测
            noise_pred = unet_pass(pipe, approximated_z_tp1, t, prompt_embeds_in, added_cond_kwargs_in)

            # 如果启用了噪声正则化且是第一步，分割批次
            if pipe.cfg.noise_regularization_num_reg_steps > 0 and i == 0:
                # 分割为"最优"噪声和当前噪声预测
                noise_pred_optimal, noise_pred = noise_pred.chunk(2)
                
                # 如果使用分类器自由引导，应用引导尺度
                if pipe.do_classifier_free_guidance:
                    noise_pred_optimal_uncond, noise_pred_optimal_text = noise_pred_optimal.chunk(2)
                    noise_pred_optimal = noise_pred_optimal_uncond + pipe.guidance_scale * (noise_pred_optimal_text - noise_pred_optimal_uncond)
                
                # 分离"最优"噪声参考
                noise_pred_optimal = noise_pred_optimal.detach()

            # 执行分类器自由引导
            if pipe.do_classifier_free_guidance:
                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                noise_pred = noise_pred_uncond + pipe.guidance_scale * (noise_pred_text - noise_pred_uncond)

            # 计算平均噪声预测（在指定范围内）
            if i >= avg_range[0] and i < avg_range[1]:
                j = i - avg_range[0]
                if nosie_pred_avg is None:
                    nosie_pred_avg = noise_pred.clone()
                else:
                    # 累积平均
                    nosie_pred_avg = j * nosie_pred_avg / (j + 1) + noise_pred / (j + 1)

        # 对噪声预测应用正则化（在指定范围或非平均模式下）
        if i >= avg_range[0] or (not pipe.cfg.average_latent_estimations and i > 0):
            noise_pred = noise_regularization(
                noise_pred, 
                noise_pred_optimal, 
                lambda_kl=pipe.cfg.noise_regularization_lambda_kl, 
                lambda_ac=pipe.cfg.noise_regularization_lambda_ac, 
                num_reg_steps=pipe.cfg.noise_regularization_num_reg_steps, 
                num_ac_rolls=pipe.cfg.noise_regularization_num_ac_rolls, 
                generator=generator
            )
        
        # 使用调度器执行反转步骤
        approximated_z_tp1 = pipe.scheduler.inv_step(noise_pred, t, z_t, **extra_step_kwargs, return_dict=False)[0].detach()

    # 如果启用了潜在估计平均，使用平均噪声执行额外步骤
    if pipe.cfg.average_latent_estimations and nosie_pred_avg is not None:
        nosie_pred_avg = noise_regularization(
            nosie_pred_avg, 
            noise_pred_optimal, 
            lambda_kl=pipe.cfg.noise_regularization_lambda_kl, 
            lambda_ac=pipe.cfg.noise_regularization_lambda_ac, 
            num_reg_steps=pipe.cfg.noise_regularization_num_reg_steps, 
            num_ac_rolls=pipe.cfg.noise_regularization_num_ac_rolls, 
            generator=generator
        )
        approximated_z_tp1 = pipe.scheduler.inv_step(nosie_pred_avg, t, z_t, **extra_step_kwargs, return_dict=False)[0].detach()

    # 执行噪声校正（如果启用）
    if pipe.cfg.perform_noise_correction:
        # 再次通过UNet获取噪声预测
        noise_pred = unet_pass(pipe, approximated_z_tp1, t, prompt_embeds, added_cond_kwargs)

        # 执行分类器自由引导
        if pipe.do_classifier_free_guidance:
            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
            noise_pred = noise_pred_uncond + pipe.guidance_scale * (noise_pred_text - noise_pred_uncond)
        
        # 使用调度器的特殊步骤更新噪声
        pipe.scheduler.step_and_update_noise(noise_pred, t, approximated_z_tp1, z_t, return_dict=False, optimize_epsilon_type=pipe.cfg.perform_noise_correction)

    return approximated_z_tp1


@torch.no_grad()
def unet_pass(pipe, z_t, t, prompt_embeds, added_cond_kwargs):
    """
    辅助函数：通过UNet进行前向传递，处理分类器自由引导和输入缩放。
    """
    # 如果需要分类器自由引导，复制输入
    latent_model_input = torch.cat([z_t] * 2) if pipe.do_classifier_free_guidance else z_t
    # 根据调度器缩放模型输入
    latent_model_input = pipe.scheduler.scale_model_input(latent_model_input, t)
    
    # 通过UNet并返回输出
    return pipe.unet(
        latent_model_input,
        t,
        encoder_hidden_states=prompt_embeds,
        timestep_cond=None,
        cross_attention_kwargs=pipe.cross_attention_kwargs,
        added_cond_kwargs=added_cond_kwargs,
        return_dict=False,
    )[0]